#include "TimeSteppers.h"

#include "Factory.h"
#include "ODESystem.h"
// #include "Infrastructure/BasicTools/include/SGLogger.h"

#include <iostream>

using namespace SG::Algebra;

//=============================================================================
// SecondOrderTimeStepper

SecondOrderTimeStepper::SecondOrderTimeStepper ()
{}

SecondOrderTimeStepper::~SecondOrderTimeStepper ()
{}

void SecondOrderTimeStepper::Initialize (SecondOrderODESystem *odes)
{
    // SecondOrderODESolver::Initialize(odes);
    this->mODEs = odes;

#if defined(SG_USE_MPI)
    void *solver = Base::Factory::GetInstance ()->Produce ("TPetscKsp<Real_t>");
#else
    void *solver = Base::Factory::GetInstance ()->Produce ("TPardiso<Real_t>");
#endif
    mSolver.reset (reinterpret_cast<LinearSolver *> (solver));
}

//=============================================================================
// CentralDifference

CentralDifference::CentralDifference ()
{}

CentralDifference::~CentralDifference ()
{}

void CentralDifference::Step (Vec_t &U, Vec_t &V, Vec_t &A)
{
    // from the k-th step to the next step
    // from time t to time t + dt

    Real_t dt = mTimeStep;
    Real_t t  = mCurrentTime;

    Real_t p0 = 1.0 / (dt * dt);
    Real_t p1 = 1.0 / (2.0 * dt);
    Real_t p2 = 2.0 * p0;
    Real_t p3 = 1.0 / p2;
    Real_t p4 = dt;

    Vec_t U_k = U;
    Vec_t V_k = V;
    Vec_t A_k = A;

    SpMat_t M, C, K;
    mODEs->ComputeImplicitJacobian (t, U_k, V_k, A_k, K, C, M);

    Vec_t F_k;
    // take U_{k} and V_{k} to evaluate F_{k}
    mODEs->ComputeExplicitValue (t, U_k, V_k, F_k);

    if (!mFactorizationIsOk)
    {
        mFactorizationIsOk = true;

        // the initial acceleration
        mODEs->ComputeDerivatives (t, U_k, V_k, A_k);

        SpMat_t K_eff = p0 * M + p1 * C;

        mSolver->compute (K_eff);

        U_km1 = U_k - p4 * V_k + p3 * A_k;
    }

    Vec_t F_m   = M * (p2 * U_k - p0 * U_km1);
    Vec_t F_c   = C * (p1 * U_km1);
    Vec_t F_u   = K * U_k;
    Vec_t F_eff = F_k + F_m + F_c - F_u;

    // U = mSolver->solve (F_eff);
    mSolver->solve (F_eff, U);

    V = p1 * (U - U_km1);
    A = p0 * (U - 2.0 * U_k + U_km1);

    // update the cached variables
    mCurrentTime = t + dt;
    U_km1        = U_k;
}

//=============================================================================
// ModifiedCentralDifference

ModifiedCentralDifference::ModifiedCentralDifference ()
{}

ModifiedCentralDifference::~ModifiedCentralDifference ()
{}

void ModifiedCentralDifference::Step (Vec_t &U, Vec_t &V, Vec_t &A)
{
    // from the k-th step to the next step
    // from time t to time t + dt

    Real_t dt = mTimeStep;
    Real_t t  = mCurrentTime;

    Real_t p0 = 1.0 / (dt * dt);
    Real_t p1 = 0.5 / dt;
    Real_t p2 = 2.0 * p0;
    //Real_t p3 = 1.0 / p2;
    Real_t p4 = dt;

    Vec_t U_k = U;
    Vec_t V_k = V;
    Vec_t A_k = A;

    SpMat_t M, C, K;
    mODEs->ComputeImplicitJacobian (t, U_k, V_k, A_k, K, C, M);

    Vec_t F_kp1;
    // take U_{k} and V_{k} to estimate F_{k+1}
    mODEs->ComputeExplicitValue (t, U_k, V_k, F_kp1);

    if (!mFactorizationIsOk)
    {
        mFactorizationIsOk = true;

        // the initial acceleration
        mODEs->ComputeDerivatives (t - dt, U_k, V_k, A_k);

        SpMat_t K_eff = p0 * M + p1 * C + 1.0 / 3.0 * K;
        mSolver->compute (K_eff);

        U_km1 = U_k - p4 * V_k;
        F_km1 = K * U_km1 + C * V_k;
        F_k   = K * U_k + C * V_k;
    }

    Vec_t F_m   = M * (p2 * U_k - p0 * U_km1);
    Vec_t F_c   = C * (p1 * U_km1);
    Vec_t F_u   = K * ((U_k + U_km1) / 3.0);

    Vec_t F_eff = 1.0 / 3.0 * (F_kp1 + F_k + F_km1) + F_m + F_c - F_u;

    // U = mSolver->solve (F_eff);
    mSolver->solve (F_eff, U);

    V = p1 * (U - U_km1);
    A = p0 * (U - 2.0 * U_k + U_km1);

    // update the cached variables
    mCurrentTime = t + dt;
    U_km1        = U_k;
    F_km1        = F_k;
    F_k          = F_kp1;
}

//=============================================================================
// NewmarkBeta

NewmarkBeta::NewmarkBeta ()
{
    // see mfem/linalg/ode.hpp:NewmarkSolver
    gamma = 0.5;
    beta  = 0.25;
}

NewmarkBeta::~NewmarkBeta ()
{}

// void Newmark::Step (const SpMat_t &M, const SpMat_t &C, const SpMat_t &K, const Vec_t &F, Vec_t &U, Vec_t &V, Vec_t &A)
// {
//     Real_t t = mCurrentTime;
//     Real_t dt = mTimeStep;

//     Real_t p0 = gamma / (beta * dt);
//     Real_t p1 = 1.0 - gamma / beta;
//     Real_t p2 = (1.0 - 0.5 * gamma / beta) * dt;
//     Real_t p3 = 1.0 / (beta * dt * dt);
//     Real_t p4 = -1.0 / (beta * dt);
//     Real_t p5 = 1 - 0.5 / beta;

//     Vec_t U_k = U;
//     Vec_t V_k = V;
//     Vec_t A_k = A;

//     Vec_t F_kp1 = F;
//     Vec_t F_m   = M * (p3 * U - p4 * V - p5 * A);
//     Vec_t F_c   = C * (p0 * U - p1 * V - p2 * A);
//     Vec_t F_eff = F_kp1 + F_m + F_c;

//     if (!mFactorizationIsOk)
//     {
//         SpMat_t K_eff = p3 * M + p0 * C + K;
//         mSolver->compute (K_eff);
//         mFactorizationIsOk = true;
//     }

//     U = mSolver->solve (F_eff);

//     V = p0 * (U - U_k) + p1 * V_k + p2 * A_k;
//     A = p3 * (U - U_k) + p4 * V_k + p5 * A_k;

//     mCurrentTime += dt;
// }

void NewmarkBeta::Step (Vec_t &U, Vec_t &V, Vec_t &A)
{
    // from the k-th step to the next step
    // from time t to time t + dt

    Real_t dt = mTimeStep;
    Real_t t  = mCurrentTime;

    // int n = mODEs->GetEquationCount ();
    Real_t beta_dt = beta * dt;

    Real_t p0 = gamma / beta_dt;
    Real_t p1 = 1.0 - gamma / beta;
    Real_t p2 = (1.0 - 0.5 * gamma / beta) * dt;
    Real_t p3 = 1.0 / (beta_dt * dt);
    Real_t p4 = -1.0 / beta_dt;
    Real_t p5 = 1 - 0.5 / beta;

    Vec_t U_k = U;
    Vec_t V_k = V;
    Vec_t A_k = A;

    SpMat_t M, C, K;
    mODEs->ComputeImplicitJacobian (t, U_k, V_k, A_k, K, C, M);

    Vec_t F_kp1;
    // take U_{k} and V_{k} to estimate F_{k+1}
    mODEs->ComputeExplicitValue (t, U_k, V_k, F_kp1);

    if (!mFactorizationIsOk)
    {
        mFactorizationIsOk = true;

        // the initial acceleration
        mODEs->ComputeDerivatives (t - dt, U_k, V_k, A_k);

        SpMat_t K_eff = p3 * M + p0 * C + K;
        mSolver->compute (K_eff);
    }

    Vec_t F_m   = M * (p3 * U_k - p4 * V_k - p5 * A_k);
    Vec_t F_c   = C * (p0 * U_k - p1 * V_k - p2 * A_k);
    Vec_t F_eff = F_kp1 + F_m + F_c;

    // U = mSolver->solve (F_eff);
    mSolver->solve (F_eff, U);

    V = p0 * (U - U_k) + p1 * V_k + p2 * A_k;
    A = p3 * (U - U_k) + p4 * V_k + p5 * A_k;

    // update the cached variables
    mCurrentTime = t + dt;
}

//=============================================================================
// WilsonTheta

WilsonTheta::WilsonTheta ()
{
    theta = 1.4;
}

WilsonTheta::~WilsonTheta ()
{}

void WilsonTheta::Step (Vec_t &U, Vec_t &V, Vec_t &A)
{
    // from the k-th step to the next step
    // from time t to time t + dt

    Real_t dt = mTimeStep;
    Real_t t  = mCurrentTime;

    Real_t theta2   = theta * theta;
    Real_t theta3   = theta2 * theta;
    Real_t theta_dt = theta * dt;

    Real_t p0  = 6.0 / (theta_dt * theta_dt);
    Real_t p1  = 3.0 / theta_dt;
    Real_t p2  = 2.0 * p1;
    Real_t p3  = theta_dt / 2.0;
    Real_t p4  = theta;
    Real_t p5  = 1.0 - theta;
    Real_t p6  = p0 / theta;
    Real_t p7  = p2 / theta;
    Real_t p8  = 1.0 - 3.0 / theta;
    Real_t p9  = 3.0 / (theta3 * dt);
    Real_t p10 = 1.0 - 3.0 / theta2;
    Real_t p11 = (1.0 - 1.5 / theta) * dt;
    Real_t p12 = 1.0 / theta3;
    Real_t p13 = 1.0 - 1.0 / theta3;
    Real_t p14 = (1.0 - 1.0 / theta2) * dt;
    Real_t p15 = 0.5 * (1.0 - 1.0 / theta) * dt * dt;

    Vec_t U_k = U;
    Vec_t V_k = V;
    Vec_t A_k = A;

    SpMat_t M, C, K;
    mODEs->ComputeImplicitJacobian (t, U_k, V_k, A_k, K, C, M);

    Vec_t F_kp1;
    // take U_{k} and V_{k} to estimate F_{k+1}
    mODEs->ComputeExplicitValue (t, U, V, F_kp1);

    if (!mFactorizationIsOk)
    {
        mFactorizationIsOk = true;

        // the initial acceleration
        mODEs->ComputeDerivatives (t - dt, U_k, V_k, A_k);

        SpMat_t K_eff = p0 * M + p1 * C + K;
        mSolver->compute (K_eff);

        // take U_{k} and V_{k} to evaluate F_{k}
        mODEs->ComputeExplicitValue (t, U_k, V_k, F_k);
    }

    Vec_t F_m = M * (p0 * U_k + p2 * V_k + 2.0 * A_k);
    Vec_t F_c = C * (p1 * U_k + 2.0 * V_k + p3 * A_k);

    Vec_t F_eff = p4 * F_kp1 + p5 * F_k + F_m + F_c;

    // Vec_t U_t = mSolver->solve (F_eff);
    Vec_t U_t;
    mSolver->solve (F_eff, U_t);

    U = p12 * U_t + p13 * U_k + p14 * V_k + p15 * A_k;
    V = p9 * (U_t - U_k) + p10 * V_k + p11 * A_k;
    A = p6 * (U_t - U_k) - p7 * V_k + p8 * A_k;

    // update the cached variables
    mCurrentTime = t + dt;
    F_k          = F_kp1;
}

//=============================================================================
// Generalized-Alpha

GeneralizedAlpha::GeneralizedAlpha ()
{
    // see mfem/linalg/ode.hpp:GeneralizedAlpha2Solver
    Real_t rho_inf = 1.0;

    rho_inf = (rho_inf > 1.0) ? 1.0 : rho_inf;
    rho_inf = (rho_inf < 0.0) ? 0.0 : rho_inf;

    alpha_m = (2.0 - rho_inf) / (1.0 + rho_inf);
    alpha_f = 1.0 / (1.0 + rho_inf);
    beta    = 0.25 * pow (1.0 + alpha_m - alpha_f, 2);
    gamma   = 0.5 + alpha_m - alpha_f;
}

GeneralizedAlpha::~GeneralizedAlpha ()
{}

void GeneralizedAlpha::Step (Vec_t &U, Vec_t &V, Vec_t &A)
{
    // from the k-th step to the next step
    // from time t to time t + dt

    Real_t dt = mTimeStep;
    Real_t t  = mCurrentTime;

    Real_t beta_dt = beta * dt;

    Real_t p0      = gamma / beta_dt;
    Real_t p1      = 1.0 - gamma / beta;
    Real_t p2      = (1.0 - 0.5 * gamma / beta) * dt;
    Real_t p3      = 1.0 / (beta_dt * dt);
    Real_t p4      = -1.0 / beta_dt;
    Real_t p5      = 1.0 - 0.5 / beta;
    Real_t p6      = 1.0 - alpha_m;
    Real_t p7      = 1.0 - alpha_f;
    Real_t p8      = alpha_f;
    Real_t p9      = p6 * p3;
    Real_t p10     = p7 * p0;
    Real_t p11     = -p6 * p4;
    Real_t p12     = 0.5 * p6 / beta - 1.0;
    Real_t p13     = p10 * dt - 1.0;
    Real_t p14     = p7 * p2;

    Vec_t U_k = U;
    Vec_t V_k = V;
    Vec_t A_k = A;

    SpMat_t M, C, K;
    mODEs->ComputeImplicitJacobian (t, U_k, V_k, A_k, K, C, M);
    
    Vec_t F_kp1;
    // take U_{k} and V_{k} to estimate F_{k+1}
    mODEs->ComputeExplicitValue (t, U_k, V_k, F_kp1);

    if (!mFactorizationIsOk)
    {
        mFactorizationIsOk = true;

        // the initial acceleration
        mODEs->ComputeDerivatives (t - dt, U_k, V_k, A_k);

        SpMat_t K_eff = p9 * M + p10 * C + p7*K;
        mSolver->compute (K_eff);

        // take U_{k} and V_{k} to evaluate F_{k}
        mODEs->ComputeExplicitValue (t, U_k, V_k, F_k);
    }

    Vec_t F_m   = M * (p9 * U_k + p11 * V_k + p12 * A_k);
    Vec_t F_c   = C * (p10 * U_k + p13 * V_k + p14 * A_k);
    Vec_t F_u   = K * (p8 * U_k);

    Vec_t F_eff = p7 * F_kp1 + p8 * F_k + F_m + F_c - F_u;

    // U = mSolver->solve (F_eff);
    mSolver->solve (F_eff, U);

    V = p0 * (U - U_k) + p1 * V_k + p2 * A_k;
    A = p3 * (U - U_k) + p4 * V_k + p5 * A_k;

    // update the cached variables
    mCurrentTime = t + dt;
    F_k          = F_kp1;
}

//=============================================================================
// Bathe

Bathe::Bathe ()
{
    gamma = 0.5;
}

Bathe::~Bathe ()
{}

void Bathe::Initialize (SecondOrderODESystem *odes)
{
    SecondOrderTimeStepper::Initialize (odes);

#if defined(SG_USE_MPI)
    void *solver = Base::Factory::GetInstance ()->Produce ("TPetscKsp<Real_t>");
#else
    void *solver = Base::Factory::GetInstance ()->Produce ("TPardiso<Real_t>");
#endif

    mSolver2.reset (reinterpret_cast<LinearSolver *> (solver));
}

void Bathe::Step (Vec_t &U, Vec_t &V, Vec_t &A)
{
    // from the k-th step to the next step
    // from time t to time t + dt

    Real_t dt = mTimeStep;
    Real_t t  = mCurrentTime;

    Real_t gamma_dt = gamma * dt;

    Real_t p0 = 2.0 / gamma_dt;
    Real_t p1 = 4.0 / gamma_dt;
    Real_t p2 = 4.0 / (gamma_dt * gamma_dt);
    Real_t p3 = gamma;
    Real_t p4 = 1.0 - gamma;
    Real_t p5 = (1.0 - gamma) / gamma_dt;
    Real_t p6 = -1.0 / (gamma_dt - gamma * gamma_dt);
    Real_t p7 = (2.0 - gamma) / (dt - gamma_dt);
    Real_t p77 = p7 * p7;
    Real_t p57 = p5 * p7;
    Real_t p67 = p6 * p7;

    // // take explicit euler to correct chen's formulations
    // Real_t p8 = p7 / dt;

    Vec_t U_k = U;
    Vec_t V_k = V;
    Vec_t A_k = A;

    SpMat_t M, C, K;
    mODEs->ComputeImplicitJacobian (t, U_k, V_k, A_k, K, C, M);

    Vec_t F_kp1;
    // take U_{k} and V_{k} to estimate F_{k+1}
    mODEs->ComputeExplicitValue (t, U_k, V_k, F_kp1);

    if (!mFactorizationIsOk)
    {
        mFactorizationIsOk = true;

        // the initial acceleration
        mODEs->ComputeDerivatives (t - dt, U_k, V_k, A_k);

        SpMat_t K_eff = p2 * M + p0 * C + K;
        mSolver->compute (K_eff);

        // SpMat_t K_eff2 = p7 * M + p7 * C + K;
        SpMat_t K_eff2 = p77 * M + p7 * C + K;
        mSolver2->compute (K_eff2);

        // take U_{k} and V_{k} to evaluate F_{k}
        mODEs->ComputeExplicitValue (t, U_k, V_k, F_k);
    }

    // t_{k+y}
    Vec_t F_m = M * (p2 * U_k + p1 * V_k + A_k);
    Vec_t F_c = C * (p0 * U_k + V_k);

    Vec_t F_eff = p3 * F_kp1 + p4 * F_k + F_m + F_c;

    // Vec_t U_kpy = mSolver->solve (F_eff);
    Vec_t U_kpy;
    mSolver->solve (F_eff, U_kpy);
    Vec_t V_kpy = p0 * (U_kpy - U_k) - V_k;

    // t_{k+1}
    // F_m = M * (p5 * V_k + p6 * V_kpy);
    F_m = M * (p5 * V_k + p6 * V_kpy + p57 * U_k + p67 * U_kpy);
    F_c = C * (p5 * U_k + p6 * U_kpy);

    F_eff = F_kp1 - F_m - F_c;

    // U = mSolver2->solve (F_eff);
    mSolver2->solve (F_eff, U);

    V = p5 * U_k + p6 * U_kpy + p7 * U;
    A = p5 * V_k + p6 * V_kpy + p7 * V;

    // update the cached variables
    mCurrentTime = t + dt;
    F_k          = F_kp1;
}
